import os, sys
from urllib.parse import quote_plus
from urllib.request import urlopen
from urllib.request import Request
from urllib.parse import urlencode
from urllib.error import URLError
try:
    import xlsxwriter
except:
    os.system("pip3 install xlsxwriter")
    import xlsxwriter
import base64
import json
import ssl
import io

ssl._create_default_https_context = ssl._create_unverified_context

def fetch_token():
    params = {'grant_type': 'client_credentials',
              'client_id': main.API_KEY,
              'client_secret': main.SECRET_KEY}
    post_data = urlencode(params)
    post_data = post_data.encode('utf-8')
    req = Request(main.TOKEN_URL, post_data)
    try:
        f = urlopen(req, timeout=5)
        result_str = f.read()
    except URLError as err:
        print(err)
    result_str = result_str.decode()

    result = json.loads(result_str)

    if ('access_token' in result.keys() and 'scope' in result.keys()):
        if not 'brain_all_scope' in result['scope'].split(' '):
            print ('please ensure has check the  ability')
            exit()
        return result['access_token']
    else:
        print ('please overwrite the correct main.API_KEY and main.SECRET_KEY')
        exit()

def read_file(image_path):
    f = None
    try:
        f = open(image_path, 'rb')
        return f.read()
    except:
        print('read image file fail')
        return None
    finally:
        if f:
            f.close()

def request(url, data):
    req = Request(url, data.encode('utf-8'))
    has_error = False
    try:
        f = urlopen(req)
        result_str = f.read()
        result_str = result_str.decode()
        return result_str
    except  URLError as err:
        print(err)

def main(a):

    image_url = main.OCR_URL + "?access_token=" + main.token
    text = ""
    file_read = read_file(a)
    result = request(image_url, urlencode({'image': base64.b64encode(file_read)}))
    result_dict = json.loads(result)
    #open(os.path.basename(a).rsplit(".")[0]+".txt", 'w', encoding='utf8').write(result)
    #open(os.path.basename(a).rsplit(".")[0]+".txt", 'w', encoding='utf8').write(str(result_dict))
    
    v_list= []
    for i in range(len(result_dict['words_result'])):
        if i==0 or result_dict['words_result'][i-1]['location']['top'] + result_dict['words_result'][i-1]['location']['height'] + result_dict['words_result'][i]['location']['height'] < result_dict['words_result'][i]['location']['top']:
            v_list.append(result_dict['words_result'][i]['words'])
        else:
            v_list[-1] = v_list[-1]+result_dict['words_result'][i]['words']

    return v_list

def main2(a):
    
    v_list=[]
    for i in a:
        i= i.strip('\'\"\\').rstrip("/")
        if os.path.isdir(i):
            for i2 in os.listdir(i):
                print(" reach : %s" %os.path.join(i, i2))
                v_list.append(main(os.path.join(i, i2)))
        else:
            print(" reach : %s" %i)
            v_list.append(main(i))
    
    workbook=xlsxwriter.Workbook('output.xlsx')
    worksheet=workbook.add_worksheet()
    for index, i in enumerate(v_list):
        print(position(index))
        worksheet.write_column(position(index), i)
    workbook.close()
    

def position(n):
    if n==0:
        return 'A1'
    digits = []
    while n:
        digits.append(int(n % 26))
        n //= 26
    v = ""
    for i in digits[::-1]:
        v+=chr(ord('A')+i)
    v+="1"
    print(v)
    return v


if __name__ == '__main__':
    #main.OCR_URL = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate_basic" #标准
    main.OCR_URL = "https://aip.baidubce.com/rest/2.0/ocr/v1/general" #标准位置
    #main.OCR_URL = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate_basic" #高精度
    #main.OCR_URL = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate" #高精度位置
    main.TOKEN_URL = 'https://aip.baidubce.com/oauth/2.0/token'

    try:
        main.API_KEY = os.environ['API_KEY_3gy6']
        main.SECRET_KEY = os.environ['SECRET_KEY_3gy6']
        main.token = fetch_token()
        is_env = True
    except:
        print("请设置环境变量 API_KEY_3gy6 和 SECRET_KEY_3gy6")
        print("https://console.bce.baidu.com/ai/?fromai=1#/ai/ocr/app/list")
        sys.exit(0)

    if len(sys.argv) <=1:
        print("请依次添加（多个）图片的路径作为命令参数\n\t每一张图片视作一竖栏，可用图片切片工具竖向切片\n\t文字的分行将自动完成")
        sys.exit(0)
    
    main2(sys.argv[1:])

